Squares of a sorted array

Time: O(N); Space: O(1); easy

Given an array of integers A sorted in non-decreasing order, return an array of the squares of each number, also in sorted non-decreasing order.

Example 1:

Input: A = [-4,-1,0,3,10]

Output: [0,1,9,16,100]

Example 2:

Input: A = [-7,-3,2,3,11]

Output: [4,9,9,49,121]

Constraints:

  • 1 <= len(A) <= 10000

  • -10000 <= A[i] <= 10000

  • A is sorted in non-decreasing order.

1. Two Pointer + Bisect [O(N), O(1)]

[4]:

import bisect

class Solution1(object):
    """
    Time: O(N)
    Space: O(1)
    """
    def sortedSquares(self, A):
        """
        :type A: List[int]
        :rtype: List[int]
        """
        right = bisect.bisect_left(A, 0)

        left = right-1
        result = []

        while 0 <= left or right < len(A):
            if right == len(A) or (0 <= left and A[left]**2 < A[right]**2):
                result.append(A[left]**2)
                left -= 1
            else:
                result.append(A[right]**2)
                right += 1

        return result

[5]:

s = Solution1()

A = [-4,-1,0,3,10]
assert s.sortedSquares(A) == [0,1,9,16,100]

A = [-7,-3,2,3,11]
assert s.sortedSquares(A) == [4,9,9,49,121]

2. Two Pointer [O(N), O(N)]

Intuition

Since the array A is sorted, loosely speaking it has some negative elements with squares in decreasing order, then some non-negative elements with squares in increasing order.

For example, with [-3, -2, -1, 4, 5, 6], we have the negative part [-3, -2, -1] with squares [9, 4, 1], and the positive part [4, 5, 6] with squares [16, 25, 36]. Our strategy is to iterate over the negative part in reverse, and the positive part in the forward direction.

Algorithm

We can use two pointers to read the positive and negative parts of the array - one pointer j in the positive direction, and another i in the negative direction.

Now that we are reading two increasing arrays (the squares of the elements), we can merge these arrays together using a two-pointer technique.

[6]:

class Solution2(object):
    """
    Time: O(N), where N is the length of A.
    Space: O(N).
    """
    def sortedSquares(self, A):
        N = len(A)
        # i, j: negative, positive parts
        j = 0
        while j < N and A[j] < 0:
            j += 1
        i = j - 1

        ans = []
        while 0 <= i and j < N:
            if A[i]**2 < A[j]**2:
                ans.append(A[i]**2)
                i -= 1
            else:
                ans.append(A[j]**2)
                j += 1

        while i >= 0:
            ans.append(A[i]**2)
            i -= 1
        while j < N:
            ans.append(A[j]**2)
            j += 1

        return ans

[7]:

s = Solution2()

A = [-4,-1,0,3,10]
assert s.sortedSquares(A) == [0,1,9,16,100]

A = [-7,-3,2,3,11]
assert s.sortedSquares(A) == [4,9,9,49,121]

3. Sort [O(NLogN), O(N)]

Intuition and Algorithm

Create an array of the squares of each element, and sort them.

[8]:

class Solution3(object):
    """
    Time: O(NLogN), where N is the length of A.
    Space: O(N).
    """
    def sortedSquares(self, A):
        return sorted(x*x for x in A)

[9]:

s = Solution3()

A = [-4,-1,0,3,10]
assert s.sortedSquares(A) == [0,1,9,16,100]

A = [-7,-3,2,3,11]
assert s.sortedSquares(A) == [4,9,9,49,121]